import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.gridspec import GridSpec
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, roc_curve, auc
from sklearn.preprocessing import StandardScaler
import talib
import warnings
from scipy import stats

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


# 获取股票基本信息
def get_stock_basic():
    df = pro.stock_basic(**{
        "ts_code": "",
        "name": "",
        "exchange": "",
        "market": "",
        "is_hs": "",
        "list_status": "L",
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "symbol",
        "name",
        "area",
        "industry",
        "cnspell",
        "market",
        "list_date",
        "act_name",
        "act_ent_type"
    ])
    return df


# 获取日线数据
def get_daily_data(ts_code, start_date, end_date):
    df = pro.daily(**{
        "ts_code": ts_code,
        "trade_date": "",
        "start_date": start_date,
        "end_date": end_date,
        "offset": "",
        "limit": ""
    }, fields=[
        "ts_code",
        "trade_date",
        "open",
        "high",
        "low",
        "close",
        "pre_close",
        "change",
        "pct_chg",
        "vol",
        "amount"
    ])
    # 按日期升序排列
    df = df.sort_values('trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    return df


# 获取财务指标
def get_financial_indicators(ts_code, start_date, end_date):
    # 获取资产负债表
    balance = pro.balancesheet(ts_code=ts_code, start_date=start_date, end_date=end_date)
    # 获取利润表
    income = pro.income(ts_code=ts_code, start_date=start_date, end_date=end_date)
    # 获取现金流量表
    cashflow = pro.cashflow(ts_code=ts_code, start_date=start_date, end_date=end_date)
    # 获取财务指标
    financial = pro.fina_indicator(ts_code=ts_code, start_date=start_date, end_date=end_date)

    return balance, income, cashflow, financial


# 策略1: KDJ策略 (传统量价策略)
def kdj_strategy(df, n=9, m1=3, m2=3):
    # 计算KDJ指标
    df['low_n'] = df['low'].rolling(n).min()
    df['high_n'] = df['high'].rolling(n).max()

    # 计算RSV
    df['rsv'] = 100 * (df['close'] - df['low_n']) / (df['high_n'] - df['low_n'])

    # 计算K、D、J值
    df['k'] = df['rsv'].ewm(com=m1 - 1, adjust=False).mean()
    df['d'] = df['k'].ewm(com=m2 - 1, adjust=False).mean()
    df['j'] = 3 * df['k'] - 2 * df['d']

    # 生成交易信号
    df['kdj_signal'] = 0

    # KDJ金叉: K线上穿D线，买入信号
    df.loc[(df['k'] > df['d']) & (df['k'].shift(1) <= df['d'].shift(1)), 'kdj_signal'] = 1

    # KDJ死叉: K线下穿D线，卖出信号
    df.loc[(df['k'] < df['d']) & (df['k'].shift(1) >= df['d'].shift(1)), 'kdj_signal'] = -1

    # 计算策略收益
    df['kdj_return'] = df['pct_chg'] * df['kdj_signal'].shift(1)
    df['kdj_cum_return'] = (1 + df['kdj_return'] / 100).cumprod() - 1

    return df


# 特征工程函数
def create_features(df, lookback=20):
    """创建用于机器学习的特征"""
    features_df = df.copy()

    # 价格相关特征
    features_df['returns'] = features_df['close'].pct_change()
    features_df['log_returns'] = np.log(1 + features_df['returns'])

    # 移动平均
    for window in [5, 10, 20, 60]:
        features_df[f'sma_{window}'] = features_df['close'].rolling(window=window).mean()
        features_df[f'sma_{window}_ratio'] = features_df['close'] / features_df[f'sma_{window}']

        features_df[f'ema_{window}'] = features_df['close'].ewm(span=window).mean()
        features_df[f'ema_{window}_ratio'] = features_df['close'] / features_df[f'ema_{window}']

    # 波动率特征
    features_df['volatility_5'] = features_df['returns'].rolling(window=5).std()
    features_df['volatility_20'] = features_df['returns'].rolling(window=20).std()

    # 交易量特征
    features_df['volume_ma_5'] = features_df['vol'].rolling(window=5).mean()
    features_df['volume_ma_20'] = features_df['vol'].rolling(window=20).mean()
    features_df['volume_ratio'] = features_df['vol'] / features_df['volume_ma_20']

    # 价格动量特征
    features_df['momentum_5'] = features_df['close'].pct_change(periods=5)
    features_df['momentum_10'] = features_df['close'].pct_change(periods=10)
    features_df['momentum_20'] = features_df['close'].pct_change(periods=20)

    # 计算技术指标
    # MACD
    features_df['macd'], features_df['macd_signal'], features_df['macd_hist'] = talib.MACD(
        features_df['close'], fastperiod=12, slowperiod=26, signalperiod=9)

    # RSI
    features_df['rsi_14'] = talib.RSI(features_df['close'], timeperiod=14)

    # 布林带
    features_df['upper_band'], features_df['middle_band'], features_df['lower_band'] = talib.BBANDS(
        features_df['close'], timeperiod=20)
    features_df['bb_width'] = (features_df['upper_band'] - features_df['lower_band']) / features_df['middle_band']

    # KDJ指标
    features_df['k'], features_df['d'] = talib.STOCH(
        features_df['high'], features_df['low'], features_df['close'],
        fastk_period=9, slowk_period=3, slowk_matype=0, slowd_period=3, slowd_matype=0)
    features_df['j'] = 3 * features_df['k'] - 2 * features_df['d']

    # 计算未来N天的收益率作为标签
    future_return_days = 5
    features_df['future_return'] = features_df['close'].pct_change(future_return_days).shift(-future_return_days)

    # 创建分类标签：1表示上涨，0表示下跌
    features_df['target'] = (features_df['future_return'] > 0).astype(int)

    # 删除NaN值
    features_df = features_df.dropna()

    return features_df


# 策略2: 基于机器学习的量化策略
def ml_strategy(df):
    # 创建特征
    features_df = create_features(df)

    # 选择特征和目标变量
    feature_columns = [
        'sma_5_ratio', 'sma_10_ratio', 'sma_20_ratio', 'sma_60_ratio',
        'ema_5_ratio', 'ema_10_ratio', 'ema_20_ratio', 'ema_60_ratio',
        'volatility_5', 'volatility_20',
        'volume_ratio',
        'momentum_5', 'momentum_10', 'momentum_20',
        'macd', 'macd_signal', 'macd_hist',
        'rsi_14',
        'bb_width',
        'k', 'd', 'j'
    ]

    X = features_df[feature_columns]
    y = features_df['target']

    # 标准化特征
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 分割训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, shuffle=False)

    # 训练随机森林模型
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]

    # 评估模型
    accuracy = accuracy_score(y_test, y_pred)
    classification_rep = classification_report(y_test, y_pred)

    # 计算ROC曲线和AUC
    fpr, tpr, _ = roc_curve(y_test, y_proba)
    roc_auc = auc(fpr, tpr)

    # 计算特征重要性
    feature_importance = pd.DataFrame({
        'feature': feature_columns,
        'importance': model.feature_importances_
    }).sort_values('importance', ascending=False)

    # 将预测结果添加到原始数据框
    predictions = pd.Series(y_pred, index=features_df.index[-len(y_pred):])
    probabilities = pd.Series(y_proba, index=features_df.index[-len(y_proba):])

    # 合并回原始数据
    ml_df = features_df.copy()
    ml_df['ml_prediction'] = predictions
    ml_df['ml_probability'] = probabilities

    # 生成交易信号
    ml_df['ml_signal'] = 0
    ml_df.loc[ml_df['ml_prediction'] == 1, 'ml_signal'] = 1  # 预测上涨，买入
    ml_df.loc[ml_df['ml_prediction'] == 0, 'ml_signal'] = -1  # 预测下跌，卖出

    # 计算策略收益
    ml_df['ml_return'] = ml_df['pct_chg'] * ml_df['ml_signal'].shift(1)
    ml_df['ml_cum_return'] = (1 + ml_df['ml_return'] / 100).cumprod() - 1

    # 返回结果和模型评估指标
    return ml_df, {
        'accuracy': accuracy,
        'classification_report': classification_rep,
        'feature_importance': feature_importance,
        'fpr': fpr,
        'tpr': tpr,
        'roc_auc': roc_auc
    }


# 计算策略绩效指标
def calculate_performance_metrics(df, strategy_name, benchmark_name='buy_hold'):
    """计算策略绩效指标"""
    metrics = {}

    # 收益率
    total_return = df[f'{strategy_name}_cum_return'].iloc[-1] * 100
    metrics['总收益率(%)'] = total_return

    # 年化收益率
    days = (df.index[-1] - df.index[0]).days
    annual_return = ((1 + total_return / 100) ** (365 / days) - 1) * 100
    metrics['年化收益率(%)'] = annual_return

    # 最大回撤
    cum_returns = 1 + df[f'{strategy_name}_cum_return']
    running_max = cum_returns.cummax()
    drawdown = (cum_returns / running_max) - 1
    max_drawdown = drawdown.min() * 100
    metrics['最大回撤(%)'] = max_drawdown

    # 夏普比率
    risk_free_rate = 0.03  # 假设无风险利率为3%
    daily_returns = df[f'{strategy_name}_return'] / 100
    sharpe_ratio = (annual_return / 100 - risk_free_rate) / (daily_returns.std() * np.sqrt(252))
    metrics['夏普比率'] = sharpe_ratio

    # 索提诺比率
    downside_returns = daily_returns[daily_returns < 0]
    if len(downside_returns) > 0:
        sortino_ratio = (annual_return / 100 - risk_free_rate) / (downside_returns.std() * np.sqrt(252))
    else:
        sortino_ratio = float('inf')
    metrics['索提诺比率'] = sortino_ratio

    # 胜率
    if len(daily_returns) > 0:
        win_rate = len(daily_returns[daily_returns > 0]) / len(daily_returns) * 100
    else:
        win_rate = 0
    metrics['胜率(%)'] = win_rate

    # 盈亏比
    if len(daily_returns[daily_returns > 0]) > 0 and len(daily_returns[daily_returns < 0]) > 0:
        profit_factor = abs(daily_returns[daily_returns > 0].mean() / daily_returns[daily_returns < 0].mean())
    else:
        profit_factor = float('inf')
    metrics['盈亏比'] = profit_factor

    # 与基准对比
    if benchmark_name in df.columns:
        benchmark_cum_return = df[f'{benchmark_name}_cum_return'].iloc[-1] * 100
        benchmark_annual_return = ((1 + benchmark_cum_return / 100) ** (365 / days) - 1) * 100

        # 超额收益
        excess_return = annual_return - benchmark_annual_return
        metrics['超额收益(%)'] = excess_return

        # 信息比率
        daily_excess_returns = daily_returns - (df[f'{benchmark_name}_return'] / 100)
        information_ratio = excess_return / 100 / (daily_excess_returns.std() * np.sqrt(252))
        metrics['信息比率'] = information_ratio

    return metrics


# 可视化KDJ策略
def visualize_kdj_strategy(stock_name, daily_df, kdj_df):
    # 创建一个带有多个子图的画布
    fig = plt.figure(figsize=(20, 18))
    gs = GridSpec(4, 1, figure=fig)

    # 子图1: 价格走势
    ax1 = fig.add_subplot(gs[0, 0])
    ax1.plot(daily_df.index, daily_df['close'], label='收盘价', color='blue')
    ax1.set_title(f'{stock_name} 价格走势', fontsize=15)
    ax1.set_ylabel('价格', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.legend()

    # 子图2: KDJ策略
    ax2 = fig.add_subplot(gs[1, 0])
    ax2.plot(kdj_df.index, kdj_df['close'], label='收盘价', color='blue')

    # 标记买入和卖出信号
    buy_signals = kdj_df[kdj_df['kdj_signal'] == 1]
    sell_signals = kdj_df[kdj_df['kdj_signal'] == -1]
    ax2.scatter(buy_signals.index, buy_signals['close'], marker='^', color='green', s=100, label='买入信号')
    ax2.scatter(sell_signals.index, sell_signals['close'], marker='v', color='red', s=100, label='卖出信号')

    ax2.set_title(f'{stock_name} KDJ策略', fontsize=15)
    ax2.set_ylabel('价格', fontsize=12)
    ax2.grid(True, linestyle='--', alpha=0.7)
    ax2.legend()

    # 子图3: KDJ指标
    ax3 = fig.add_subplot(gs[2, 0])
    ax3.plot(kdj_df.index, kdj_df['k'], label='K线', color='blue')
    ax3.plot(kdj_df.index, kdj_df['d'], label='D线', color='red')
    ax3.plot(kdj_df.index, kdj_df['j'], label='J线', color='green')
    ax3.axhline(y=80, color='gray', linestyle='--', alpha=0.7)
    ax3.axhline(y=20, color='gray', linestyle='--', alpha=0.7)
    ax3.set_title('KDJ指标', fontsize=15)
    ax3.set_ylabel('值', fontsize=12)
    ax3.grid(True, linestyle='--', alpha=0.7)
    ax3.legend()

    # 子图4: 策略收益
    ax4 = fig.add_subplot(gs[3, 0])
    ax4.plot(kdj_df.index, kdj_df['kdj_cum_return'], label='KDJ策略收益', color='red', linewidth=2)
    ax4.plot(daily_df.index, (1 + daily_df['pct_chg'] / 100).cumprod() - 1, label='买入持有收益', color='blue',
             linewidth=2)
    ax4.set_title('KDJ策略收益对比', fontsize=15)
    ax4.set_ylabel('累计收益率', fontsize=12)
    ax4.set_xlabel('日期', fontsize=12)
    ax4.grid(True, linestyle='--', alpha=0.7)

    # 添加策略绩效指标
    kdj_return = kdj_df['kdj_cum_return'].iloc[-1] * 100
    buy_hold_return = ((1 + daily_df['pct_chg'] / 100).cumprod() - 1).iloc[-1] * 100

    ax4.text(0.05, 0.95, f'KDJ策略收益率: {kdj_return:.2f}%', transform=ax4.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='red', alpha=0.2))
    ax4.text(0.05, 0.90, f'买入持有收益率: {buy_hold_return:.2f}%', transform=ax4.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='blue', alpha=0.2))

    ax4.legend()

    plt.tight_layout()
    plt.savefig(f'{stock_name}_KDJ策略分析.png', dpi=300, bbox_inches='tight')
    plt.show()


# 可视化机器学习策略
def visualize_ml_strategy(stock_name, daily_df, ml_df, ml_metrics):
    # 创建一个带有多个子图的画布
    fig = plt.figure(figsize=(20, 18))
    gs = GridSpec(4, 1, figure=fig)

    # 子图1: 价格走势
    ax1 = fig.add_subplot(gs[0, 0])
    ax1.plot(daily_df.index, daily_df['close'], label='收盘价', color='blue')
    ax1.set_title(f'{stock_name} 价格走势', fontsize=15)
    ax1.set_ylabel('价格', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.legend()

    # 子图2: 机器学习策略
    ax2 = fig.add_subplot(gs[1, 0])
    ax2.plot(ml_df.index, ml_df['close'], label='收盘价', color='blue')

    # 标记买入和卖出信号
    buy_signals = ml_df[ml_df['ml_signal'] == 1]
    sell_signals = ml_df[ml_df['ml_signal'] == -1]
    ax2.scatter(buy_signals.index, buy_signals['close'], marker='^', color='green', s=100, label='买入信号')
    ax2.scatter(sell_signals.index, sell_signals['close'], marker='v', color='red', s=100, label='卖出信号')

    ax2.set_title(f'{stock_name} 机器学习策略', fontsize=15)
    ax2.set_ylabel('价格', fontsize=12)
    ax2.grid(True, linestyle='--', alpha=0.7)
    ax2.legend()

    # 子图3: 策略收益
    ax3 = fig.add_subplot(gs[2, 0])
    ax3.plot(ml_df.index, ml_df['ml_cum_return'], label='机器学习策略收益', color='green', linewidth=2)
    ax3.plot(daily_df.index, (1 + daily_df['pct_chg'] / 100).cumprod() - 1, label='买入持有收益', color='blue',
             linewidth=2)
    ax3.set_title('机器学习策略收益对比', fontsize=15)
    ax3.set_ylabel('累计收益率', fontsize=12)
    ax3.set_xlabel('日期', fontsize=12)
    ax3.grid(True, linestyle='--', alpha=0.7)

    # 添加策略绩效指标
    ml_return = ml_df['ml_cum_return'].iloc[-1] * 100
    buy_hold_return = ((1 + daily_df['pct_chg'] / 100).cumprod() - 1).iloc[-1] * 100

    ax3.text(0.05, 0.95, f'机器学习策略收益率: {ml_return:.2f}%', transform=ax3.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='green', alpha=0.2))
    ax3.text(0.05, 0.90, f'买入持有收益率: {buy_hold_return:.2f}%', transform=ax3.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='blue', alpha=0.2))

    ax3.legend()

    # 子图4: ROC曲线
    ax4 = fig.add_subplot(gs[3, 0])
    fpr = ml_metrics['fpr']
    tpr = ml_metrics['tpr']
    roc_auc = ml_metrics['roc_auc']

    ax4.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
    ax4.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    ax4.set_xlim([0.0, 1.0])
    ax4.set_ylim([0.0, 1.05])
    ax4.set_xlabel('假阳性率', fontsize=12)
    ax4.set_ylabel('真阳性率', fontsize=12)
    ax4.set_title('机器学习模型ROC曲线', fontsize=15)
    ax4.legend(loc="lower right")
    ax4.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig(f'{stock_name}_机器学习策略分析.png', dpi=300, bbox_inches='tight')
    plt.show()


# 可视化特征重要性
def visualize_feature_importance(stock_name, ml_metrics):
    # 创建一个带有多个子图的画布
    fig = plt.figure(figsize=(20, 15))
    gs = GridSpec(3, 1, figure=fig)

    # 子图1: 特征重要性柱状图
    ax1 = fig.add_subplot(gs[0, 0])
    feature_importance = ml_metrics['feature_importance']
    sns.barplot(x='importance', y='feature', data=feature_importance, ax=ax1)
    ax1.set_title(f'{stock_name} 机器学习模型特征重要性', fontsize=15)
    ax1.set_xlabel('重要性', fontsize=12)
    ax1.set_ylabel('特征', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)

    # 子图2: 特征重要性折线图
    ax2 = fig.add_subplot(gs[1, 0])
    ax2.plot(range(len(feature_importance)), feature_importance['importance'], marker='o', linestyle='-')
    ax2.set_title('特征重要性分布', fontsize=15)
    ax2.set_xlabel('特征排名', fontsize=12)
    ax2.set_ylabel('重要性', fontsize=12)
    ax2.set_xticks(range(len(feature_importance)))
    ax2.set_xticklabels(feature_importance['feature'], rotation=90)
    ax2.grid(True, linestyle='--', alpha=0.7)

    # 子图3: 累积特征重要性
    ax3 = fig.add_subplot(gs[2, 0])
    cumulative_importance = feature_importance['importance'].cumsum()
    ax3.bar(range(len(feature_importance)), feature_importance['importance'], alpha=0.7, label='单个特征重要性')
    ax3.plot(range(len(feature_importance)), cumulative_importance, marker='o', color='red',
             linewidth=2, label='累积重要性')
    ax3.axhline(y=0.8, color='green', linestyle='--', alpha=0.7, label='80%阈值')

    # 找出累积重要性达到80%的特征数量
    n_features_80 = np.argmax(cumulative_importance >= 0.8) + 1
    ax3.axvline(x=n_features_80 - 1, color='blue', linestyle='--', alpha=0.7,
                label=f'80%重要性需要{n_features_80}个特征')

    ax3.set_title('累积特征重要性', fontsize=15)
    ax3.set_xlabel('特征排名', fontsize=12)
    ax3.set_ylabel('重要性', fontsize=12)
    ax3.set_xticks(range(len(feature_importance)))
    ax3.set_xticklabels(feature_importance['feature'], rotation=90)
    ax3.grid(True, linestyle='--', alpha=0.7)
    ax3.legend()

    plt.tight_layout()
    plt.savefig(f'{stock_name}_特征重要性分析.png', dpi=300, bbox_inches='tight')
    plt.show()


# 可视化策略比较
def visualize_strategy_comparison(stock_name, kdj_df, ml_df):
    # 创建一个带有多个子图的画布
    fig = plt.figure(figsize=(20, 12))
    gs = GridSpec(2, 1, figure=fig)

    # 子图1: 策略收益对比
    ax1 = fig.add_subplot(gs[0, 0])
    ax1.plot(kdj_df.index, kdj_df['kdj_cum_return'], label='KDJ策略收益', color='red', linewidth=2)
    ax1.plot(ml_df.index, ml_df['ml_cum_return'], label='机器学习策略收益', color='green', linewidth=2)
    ax1.plot(kdj_df.index, (1 + kdj_df['pct_chg'] / 100).cumprod() - 1, label='买入持有收益', color='blue', linewidth=2)
    ax1.set_title('策略收益对比', fontsize=15)
    ax1.set_ylabel('累计收益率', fontsize=12)
    ax1.set_xlabel('日期', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)

    # 添加策略绩效指标
    kdj_return = kdj_df['kdj_cum_return'].iloc[-1] * 100
    ml_return = ml_df['ml_cum_return'].iloc[-1] * 100
    buy_hold_return = ((1 + kdj_df['pct_chg'] / 100).cumprod() - 1).iloc[-1] * 100

    ax1.text(0.05, 0.95, f'KDJ策略收益率: {kdj_return:.2f}%', transform=ax1.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='red', alpha=0.2))
    ax1.text(0.05, 0.90, f'机器学习策略收益率: {ml_return:.2f}%', transform=ax1.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='green', alpha=0.2))
    ax1.text(0.05, 0.85, f'买入持有收益率: {buy_hold_return:.2f}%', transform=ax1.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='blue', alpha=0.2))

    ax1.legend()

    # 计算策略绩效指标
    kdj_metrics = calculate_performance_metrics(kdj_df, 'kdj')
    ml_metrics = calculate_performance_metrics(ml_df, 'ml')

    # 子图2: 策略绩效雷达图
    ax2 = fig.add_subplot(gs[1, 0], polar=True)

    # 选择要显示的指标
    radar_metrics = ['年化收益率(%)', '最大回撤(%)', '夏普比率', '胜率(%)', '盈亏比']

    # 转换指标值为0-1之间
    max_values = {
        '年化收益率(%)': max(kdj_metrics['年化收益率(%)'], ml_metrics['年化收益率(%)'], 20),
        '最大回撤(%)': 30,  # 负值，越小越好
        '夏普比率': max(kdj_metrics['夏普比率'], ml_metrics['夏普比率'], 2),
        '胜率(%)': 100,
        '盈亏比': max(kdj_metrics['盈亏比'], ml_metrics['盈亏比'], 3)
    }

    # 准备雷达图数据
    kdj_radar = [
        kdj_metrics['年化收益率(%)'] / max_values['年化收益率(%)'],
        1 - abs(kdj_metrics['最大回撤(%)'] / max_values['最大回撤(%)']),  # 最大回撤取反，因为越小越好
        kdj_metrics['夏普比率'] / max_values['夏普比率'],
        kdj_metrics['胜率(%)'] / max_values['胜率(%)'],
        kdj_metrics['盈亏比'] / max_values['盈亏比']
    ]

    ml_radar = [
        ml_metrics['年化收益率(%)'] / max_values['年化收益率(%)'],
        1 - abs(ml_metrics['最大回撤(%)'] / max_values['最大回撤(%)']),  # 最大回撤取反，因为越小越好
        ml_metrics['夏普比率'] / max_values['夏普比率'],
        ml_metrics['胜率(%)'] / max_values['胜率(%)'],
        ml_metrics['盈亏比'] / max_values['盈亏比']
    ]

    # 角度
    angles = np.linspace(0, 2 * np.pi, len(radar_metrics), endpoint=False).tolist()
    angles += angles[:1]  # 闭合雷达图

    # 添加数据点
    kdj_radar += kdj_radar[:1]
    ml_radar += ml_radar[:1]

    # 绘制雷达图
    ax2.plot(angles, kdj_radar, 'o-', linewidth=2, label='KDJ策略', color='red')
    ax2.plot(angles, ml_radar, 'o-', linewidth=2, label='机器学习策略', color='green')
    ax2.fill(angles, kdj_radar, alpha=0.25, color='red')
    ax2.fill(angles, ml_radar, alpha=0.25, color='green')

    # 设置标签
    ax2.set_thetagrids(np.degrees(angles[:-1]), radar_metrics)
    ax2.set_title('策略绩效对比雷达图', fontsize=15)
    ax2.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1))

    plt.tight_layout()
    plt.savefig(f'{stock_name}_策略比较.png', dpi=300, bbox_inches='tight')
    plt.show()


# 主函数
def main():
    # 获取股票基本信息
    stocks = get_stock_basic()

    # 选择几支股票进行分析
    selected_stocks = stocks[(stocks['industry'].isin(['银行', '白酒', '电子'])) &
                             (stocks['market'] == '主板')].head(3)

    # 设置时间范围
    start_date = '20200101'
    end_date = '20230101'

    # 对每支股票进行
    # 对每支股票进行分析
    for _, stock in selected_stocks.iterrows():
        ts_code = stock['ts_code']
        stock_name = stock['name']
        print(f"正在分析 {stock_name} ({ts_code})...")

        # 获取日线数据
        daily_data = get_daily_data(ts_code, start_date, end_date)

        # 应用KDJ策略
        kdj_data = kdj_strategy(daily_data.copy())

        # 应用机器学习策略
        ml_data, ml_metrics = ml_strategy(daily_data.copy())

        # 可视化KDJ策略
        visualize_kdj_strategy(stock_name, daily_data, kdj_data)

        # 可视化机器学习策略
        visualize_ml_strategy(stock_name, daily_data, ml_data, ml_metrics)

        # 可视化特征重要性（单独一页）
        visualize_feature_importance(stock_name, ml_metrics)

        # 可视化策略比较
        visualize_strategy_comparison(stock_name, kdj_data, ml_data)

        # 打印策略绩效指标
        print(f"\n{stock_name} 策略绩效:")

        # 计算并打印KDJ策略绩效
        kdj_metrics = calculate_performance_metrics(kdj_data, 'kdj')
        print("\nKDJ策略:")
        for metric, value in kdj_metrics.items():
            print(f"{metric}: {value:.4f}")

        # 计算并打印机器学习策略绩效
        ml_metrics_dict = calculate_performance_metrics(ml_data, 'ml')
        print("\n机器学习策略:")
        for metric, value in ml_metrics_dict.items():
            print(f"{metric}: {value:.4f}")

        print("-" * 70)


if __name__ == "__main__":
    main()
